module Zeno.Solver
(
  solve
)
where

import Prelude ()
import StdImports
import qualified Data.Set as Set

import Zeno.Core
import Zeno.ProofLogic ( Proof (..) )
import qualified Zeno.ProofLogic as PL 
import Zeno.Debug

{- 
TODO

Test whether "contradiction" ever needs a depth greater than one
  but it makes sense that it would
  
Do things in phases?
  like after induction there is no generalisation?
  after splitting there is no induction?
  
Remember to add kterm hypotheses to the conditions
  for elem-min
  
Theory about removal of multiple linked induction hypotheses on application

Theory about only adding a hypothesis one the "first" inductive step of a chain
  that is to say do not add a hypothesis if hypotheses already exist
  
Try merging case-splitting back into induction
  just remember to apply function definitions whilst finding critical terms
  
*******
Express ALL/EXISTS in Haskell itself,
  i.e. a non-terminating function which builds up every single 
    different value of a type
  see if we can prove things about this,
    a generalised notion of induction
*******
-}

data SolverState
  = SolverState   { sZeno :: !Zeno,
                    sGoal :: !ZClause,
                    sHypotheses :: ![ZClause],
                    sDoGeneralisation :: !Bool,
                    sIndSteps :: !Int,
                    sConSteps :: !Int,
                    sRandom :: !StdGen }
    
type Solver = Reader SolverState 
type StatefulSolver = State SolverState

liftZenoReader :: ZenoReader a -> Solver a
liftZenoReader zr = do
  zeno <- sZeno <$> ask
  return $ runReader zr zeno
  
liftZenoState :: ZenoState a -> StatefulSolver a
liftZenoState zs = do
  zeno <- sZeno <$> get
  let (a, zeno') = runState zs zeno
  modify $ \s -> s { sZeno = zeno' }
  return a
  
solve :: (Int, Int) -> String -> StdGen -> ZenoReader ZProof
solve (ind, con) lemma random = do
  lemma <- getLemma lemma
  zeno <- ask
  let state = 
        SolverState { sZeno = zeno,
                      sGoal = lemma { clauseVariables = [] },
                      sHypotheses = mempty,
                      sDoGeneralisation = True,
                      sIndSteps = ind,
                      sConSteps = con,
                      sRandom = random }
  return $ runReader solve0 state

-- Applying trivial definitions
solve0 :: Solver ZProof
solve0 = do
  Clause [] eq conds <- sGoal <$> ask
  let eq' = mapPair (applyTermSubstitution (Set.fromList conds)) eq 
  (goal', changed) <- liftZenoReader $ 
    applyTermDefinitions conds (Clause [] eq' conds)
  if not changed
    then solve1
    else do
      proof <- localGoal goal' solve1
      step <- proofstep "def"
      return $ PL.addStep step proof

-- Find and apply any conditions to the equality
solve1 :: Solver ZProof
solve1 = solve2
{- do
  hyps <- sHypotheses <$> ask
  defs <- clauseConditions <$> sGoal <$> ask
  let def_sub = applyTermSubstitution (Set.fromList defs)
      applicable = do
        hyp <- hyps
        guard . constructorTerm . snd . clauseEquality $ hyp
        guard . not . null . clauseConditions $ hyp
        let hyp' = removeTrivialConds (def_sub hyp)
        guard . null . clauseConditions $ hyp'
        return (clauseEquality hyp', hyp)
  if null applicable
    then solve2
    else do
      step <- proofstep "add"
      let (new_conds, used_hyps) = unzip applicable
          solve2' = 
            removeHypotheses hyps $
            addConditions new_conds $
            PL.addStep step <$> solve2
      PL.anyProven <$> randomSequence [solve2', solve2]
  where 
    removeTrivialConds :: ZClause -> ZClause
    removeTrivialConds cls = cls 
      { clauseConditions = filter (uncurry (/=)) (clauseConditions cls) } -}
      
-- Apply inductive hypotheses
solve2 :: Solver ZProof
solve2 = checkGoal $ do
  goal <- sGoal <$> ask
  hyps <- sHypotheses <$> ask    
  apply_hyps' <- concatMapM maybeApplyHyp hyps
  step <- proofstep "hyp"
  let apply_hyps = map (PL.addStep step <$>) apply_hyps' 
  PL.anyProven <$> randomSequence (solve3 : apply_hyps)  
  where
    safeSubstitution :: (ZTerm, ZTerm) -> Bool
    safeSubstitution (t1, t2) = termType t1 == termType t2
  
    maybeApplyHyp :: ZClause -> Solver [Solver ZProof]
    maybeApplyHyp hyp = do
      (l_goal, r_goal) <- clauseEquality <$> sGoal <$> ask
      apply_left <- findApplications l_goal hyp
      apply_right <- findApplications r_goal (flipClauseEquality hyp)
      let l_eqs = map (\t -> (t, r_goal)) apply_left
          r_eqs = map (\t -> (l_goal, t)) apply_right
          all_eqs = l_eqs ++ r_eqs
      return $
        map (removeHypothesis hyp) $
        map enableGeneralisation $
        map (flip localGoalEquality solve0) all_eqs
      
    findApplications :: ZTerm -> ZClause -> Solver [ZTerm]
    findApplications goal_term 
                     (Clause vars (from_hyp, to_hyp) conds) = 
      concatMapM applyUnifier $ 
        filter validUnifier $ 
          allUnifiers from_hyp goal_term
      where
        validUnifier :: ZTermSubstitution -> Bool
        validUnifier = all (flip elem (map Var vars) . fst)
        
        applyUnifier :: ZTermSubstitution -> Solver [ZTerm]
        applyUnifier uni = do
          defs <- clauseConditions <$> sGoal <$> ask
          let sub = applyTermSubstitution uni
              (terms, kterms) = unzip conds
              defs_sub = applyTermSubstitution (Set.fromList defs)
              terms' = map (defs_sub . sub) terms
          if not $ all (uncurry (==)) (terms' `zip` kterms)
            then return []
            else do
              (l_goal, r_goal) <- clauseEquality <$> sGoal <$> ask
              let goal_vars = nub $ variables l_goal ++ variables r_goal
                  hyp_vars = nub $ variables to_hyp' ++ variables from_hyp'
                  repl_vars = goal_vars \\ hyp_vars
                  hyp_sub@(from_hyp', to_hyp') = (sub from_hyp, sub to_hyp)
                  free_vars = filter (flip elem vars) (variables to_hyp')
                  subs = sequence $ do
                    repl <- repl_vars
                    return $ do
                      var <- free_vars
                      return (Var var, Var repl)
                  subs' = filter (all safeSubstitution) subs
                  hyp_subs = do
                    new_sub <- Set.fromList <$> subs'
                    return (from_hyp', applyTermSubstitution new_sub to_hyp')
                  hyp_subs' =
                    if null hyp_subs
                      then [Set.singleton hyp_sub]
                      else map Set.singleton hyp_subs
              return $ map (flip applyTermSubstitution goal_term) hyp_subs'
              
-- Apply factoring
solve3 :: Solver ZProof
solve3 = checkGoal $ do
  goal <- sGoal <$> ask
  let (l_term,  r_term) = clauseEquality goal
      (l_func : l_args) = flattenTerm l_term
      (r_func : r_args) = flattenTerm r_term
  if l_func /= r_func
    then solve4
    else do
      step <- proofstep "fac"
      let factor_goals = l_args `zip` r_args
          --nontrivial = filter (uncurry (/=)) factor_goals
          factor_proofs = mapM (flip localGoalEquality solve3) factor_goals
          factor_proof = PL.allProvenStep step <$> factor_proofs
      if constructorTerm l_func 
        then factor_proof
        else PL.anyProven <$> randomSequence [solve4, factor_proof] 
        
-- Apply inductive, generalisation and case split steps
solve4 :: Solver ZProof
solve4 = checkGoal $ do
  goal@(Clause [] (l_term, r_term) conds) <- sGoal <$> ask
  let all_terms = l_term : r_term : (map fst conds)
      eq_vars = toList l_term ++ toList r_term
  strict_terms <- filter simpleTypeTerm . nub <$>
    mapM strictTerm all_terms
  
  let (inds, rest) = partition varTerm strict_terms
      inds' = filter variableTerm inds -- \\ all_terms
      terms_terms = concatMap termList all_terms
      (genrs, splits) = partition (flip elem terms_terms) rest
      inds'' = map goInductive (map fromVar inds')
      splits' = map goSplit (splits \\ all_terms)
      
  do_genr <- sDoGeneralisation <$> ask
  all_genrs' <- if not do_genr                            
    then return []
    else do
      global_genrs <- concatMapM globalGeneralisations (inds' ++ genrs)
      let all_genrs = global_genrs -- \\ all_terms
      return (map applyGeneralisation all_genrs)
    
  ind_steps <- sIndSteps <$> ask
  let proofs = 
        (if null splits' then inds'' else [head splits']) ++ all_genrs'
  if ind_steps > 0
    then PL.anyProven <$> randomSequence proofs
    else return Unknown
  where
    goSplit :: ZTerm -> Solver ZProof
    goSplit term = do
      let Type dtype = termType term
          cons = zDataTypeCons dtype
      splits <- stateful (mapM (liftZenoState . instantiate) cons) 
                         (mapM goSplit')
      term_s <- liftZenoReader $ showZTerm term
      step <- proofstep ("cse " ++ term_s)
      return $ PL.allProvenStep step splits
      where
        goSplit' :: (ZTerm, [ZVar]) -> Solver ZProof
        goSplit' (kterm, _) = do
          Clause [] eq conds <- sGoal <$> ask
          let goal = Clause [] eq ((term, kterm) : conds)
          disableGeneralisation (localGoal goal solve0)
    
    goInductive :: ZVar -> Solver ZProof
    goInductive var = do
      let Type dtype = zVarType var
          cons = zDataTypeCons dtype
          inds = stateful (mapM (liftZenoState . instantiate) cons) 
                          (mapM goInductive')
      var_s <- liftZenoReader $ showZVar var
      step <- proofstep ("ind " ++ var_s)
      let proof = PL.allProvenStep step <$> inds
      Clause [] eq conds <- sGoal <$> ask
      let containsVar (t1, t2) = 
            var `elem` toList t1 || var `elem` toList t2
          (conds_with, conds_without) = partition containsVar conds
          conds_pset = map (++ conds_without) (subsequences conds_with)
      let proofs = map (flip localGoalConditions proof) conds_pset
      PL.anyProven <$> randomSequence proofs
      where
        goInductive' :: (ZTerm, [ZVar]) -> Solver ZProof
        goInductive' (ind, ind_vars) = do
          goal <- sGoal <$> ask
          let ind_sub = Set.singleton (Var var, ind)
          
              neededCondition :: ZEquality -> Bool
              neededCondition (t1, t2) = 
                var `elem` toList t1 || var `elem` toList t2
          
              hyp_goal = goal 
                { clauseConditions = 
                    filter neededCondition (clauseConditions goal) }
                
              createHyp :: ZVar -> StatefulSolver ZClause
              createHyp ind_var = liftZenoState $ do
                let others = filter (/= var) (variables hyp_goal)
                others' <- mapM freshVariable others
                let sublist = (var, ind_var) : (others `zip` others')
                    sub = Set.fromList $ map (Var *** Var) sublist
                    hyp = applyTermSubstitution sub hyp_goal
                return $ hyp { clauseVariables = others' }
                
          stateful (mapM createHyp ind_vars) $ \hyps -> do
            substituteTerms ind_sub $ 
              disableGeneralisation $
              takeIndStep $
              addHypotheses hyps solve0
      
    globalGeneralisations :: ZTerm -> Solver [ZTerm]
    globalGeneralisations term = do
      Clause [] eq conds <- sGoal <$> ask
       
      let subterms :: ZEquality -> [ZTerm]
          subterms (t1, t2) = 
            filter (flip containsTerm term) $
              (termList t1 ++ termList t2)
            
          duplicated :: Eq a => [a] -> [a]
          duplicated xs = nub (xs \\ nub xs)
          
          all_subterms = 
            filter simpleTypeTerm $
            filter appTerm $ 
            duplicated $
              concatMap subterms (eq : conds)
            
          isFactor :: ZTerm -> Bool
          isFactor term = any (flip containsSubTerm term) all_subterms
 
      return $ filter (not . isFactor) all_subterms
        
    applyGeneralisation :: ZTerm -> Solver ZProof
    applyGeneralisation term = do
      stateful (liftZenoState newIdS) $ \new_id -> do
        let new_var = ZVar new_id (termType term) ZVariable
            term_sub = Set.singleton (term, Var new_var)
        goal <- sGoal <$> ask
        term_s <- liftZenoReader $ showZTerm term
        var_s <- liftZenoReader $ showZVar new_var
        step <- proofstep $ "gen " ++ term_s ++ " => " ++ var_s
        PL.addStep step <$> substituteTerms term_sub solve3
      
    strictTerm :: ZTerm -> Solver ZTerm
    strictTerm term = do
      maybe_expr <- liftZenoReader $ applyFunctionDef [] term
      case maybe_expr of
        Nothing -> 
          return term
        Just (Value {}) -> do
          term_s <- liftZenoReader (showZTerm term)
          error $ "Earlier trivial step not applied to \"" ++ term_s ++ "\""
        Just (CaseOf term' _) -> do
          if term `containsTerm` term'
            then strictTerm term'
            else return term'

contradiction :: [ZEquality] -> Int -> Solver Bool
contradiction eqs depth 
  | any (uncurry unequalTerms) eqs = return True
  | otherwise = do
    inverses <- mapM solveInverse eqs
    if any (/= Unknown) inverses
      then return True
      else return False
  where 
    solveInverse :: ZEquality -> Solver ZProof
    solveInverse (t1, t2)
      | not (constructorTerm t2) = return Unknown
      | otherwise = do
        let ((Var con) : _) = flattenTerm t2
            Type dtype = termType t2
            conterms = map Var $ filter (/= con) (zDataTypeCons dtype)
        if not (all simpleTypeTerm conterms)
          then return Unknown
          else simpleSolver $ PL.anyProven <$> 
            mapM (\t2' -> localGoal (Clause [] (t1, t2') []) solve0) conterms
 
    simpleSolver :: Solver a -> Solver a
    simpleSolver = local $ \s -> s 
      { sConSteps = 0,
        sIndSteps = depth,
        sHypotheses = mempty,
        sDoGeneralisation = True }
    
{-
  let (terms, kterms) = unzip eqs
  terms' <- map fst <$> mapM (applyTermDefinitions []) terms
  let eqs' = (terms' `zip` kterms)
  if any (uncurry unequalTerms) eqs' 
      then return True
    else if all (uncurry (==)) eqs' || depth <= 0
      then return False
    else do
      term_splits <- concatMapM splits (filter simpleTypeTerm terms')
      results <- mapM (splitOn eqs') term_splits
      return (or results)
  where
    splitOn :: [ZEquality] -> ZTerm -> ZenoReader Bool
    splitOn eqs term = do
      let Type dtype = termType term
          cons = zDataTypeCons dtype
      results <- stateful (mapM instantiate cons) (mapM testSplit)
      return (and results)
      where
        testSplit :: (ZTerm, [ZVar]) -> ZenoReader Bool
        testSplit (term', _) = do
          let sub = applyTermSubstitution (Set.singleton (term, term'))
          contradiction (map (first sub) eqs) (depth - 1)
    
    splits :: ZTerm -> ZenoReader [ZTerm]
    splits var@(Var (ZVar _ _ ZVariable)) = return [var]
    splits term = do
      maybe_expr <- applyFunctionDef [] term
      case maybe_expr of
        Nothing -> return []
        Just (Value term') -> return [term']
        Just (CaseOf term' _) -> do
          terms <- splits term'
          return (term : terms)
  -}        
invalidTerm :: ZTerm -> Bool
invalidTerm term = 
  let ((Var func) : args) = flattenTerm term
      arg_types = flattenType (zVarType func)
      res_types = drop (length args) arg_types 
  in null res_types
       
randomSequence :: [Solver a] -> Solver [a]
randomSequence solvers = do
  random <- sRandom <$> ask
  let (solvers', random') = shuffle solvers random
  localRandom random' (splitGens solvers')
  where
    localRandom :: StdGen -> Solver a -> Solver a
    localRandom gen = local (\s -> s { sRandom = gen })
        
    splitGens :: [Solver a] -> Solver [a]
    splitGens = foldl' splitGen (return [])
      where 
        splitGen :: Solver [a] -> Solver a -> Solver [a]
        splitGen ss s = do
          random <- sRandom <$> ask
          let (r1, r2) = split random
          ss' <- localRandom r2 ss
          s' <- localRandom r1 s
          return (s' : ss')
          
substituteTerms :: ZTermSubstitution -> Solver a -> Solver a
substituteTerms sub = 
  let f = applyTermSubstitution sub 
  in local $ \s -> s
      { sGoal = f (sGoal s), 
        sHypotheses = map f (sHypotheses s) }
        
enableGeneralisation :: Solver a -> Solver a
enableGeneralisation = local $ \s -> s { sDoGeneralisation = True }

disableGeneralisation :: Solver a -> Solver a
disableGeneralisation = local $ \s -> s { sDoGeneralisation = False }
    
localGoal :: ZClause -> Solver a -> Solver a
localGoal goal = local $ \s -> s { sGoal = goal }

localGoalEquality :: ZEquality -> Solver a -> Solver a
localGoalEquality eq = local $ \s -> s 
  { sGoal = (sGoal s) { clauseEquality = eq } }

localGoalConditions :: [ZEquality] -> Solver a -> Solver a
localGoalConditions conds = local $ \s -> s
  { sGoal = (sGoal s) { clauseConditions = conds } }
  
addCond :: ZEquality -> Solver a -> Solver a
addCond cond = local $ \s -> s
  { sGoal = (sGoal s) 
    { clauseConditions = cond : (clauseConditions (sGoal s)) } }
  
addConditions :: [ZEquality] -> Solver a -> Solver a
addConditions = foldl (\f c -> f . addCond c) id
    
addHypothesis :: ZClause -> Solver a -> Solver a
addHypothesis cls = local $ \s -> s
  { sHypotheses = cls : (sHypotheses s) }
  
removeHypothesis :: ZClause -> Solver a -> Solver a
removeHypothesis cls = local $ \s -> s
  { sHypotheses = delete cls (sHypotheses s) }
  
removeHypotheses :: [ZClause] -> Solver a -> Solver a
removeHypotheses = foldl (\f c -> f . removeHypothesis c) id

removeAllHypotheses :: Solver a -> Solver a
removeAllHypotheses = local $ \s -> s { sHypotheses = [] }

addHypotheses :: [ZClause] -> Solver a -> Solver a
addHypotheses = foldl (\f c -> f . addHypothesis c) id
  
debugShowGoal :: Solver a -> Solver a
debugShowGoal solver = do
  goal <- sGoal <$> ask
  shown <- liftZenoReader $ showZClause goal
  traceZ shown $ solver
  
-- Check whether the goal is trivially true or false
checkGoal :: Solver ZProof -> Solver ZProof
checkGoal solver = do
  goal@(Clause [] (t1, t2) conds) <- sGoal <$> ask
  meh <- liftZenoReader $ mapMaybeM (applyFunctionDef []) (map fst conds)
  meh2 <- liftZenoReader $ intercalate "\n" <$> mapM showZExpression meh
  meh3 <- liftZenoReader $ showZClause goal
  steps <- sConSteps <$> ask
  contr <- contradiction (clauseConditions goal) steps
  -- Check whether the conditions are contradictory
  if contr
      then return $ Proven (PL.ProofStep "con" goal) []
    else if t1 == t2
      then return $ Proven (PL.ProofStep "eql" goal) []
    else if unequalTerms t1 t2
      then return Unknown
      --then return $ Disproven ("irf", goal) []
    else do
      -- Filter out trivial conditions
      let conds' = filter (uncurry (/=)) conds 
      localGoal (Clause [] (t1, t2) conds') solver

takeIndStep :: Solver a -> Solver a 
takeIndStep = local $ \s -> 
  s { sIndSteps = (sIndSteps s) - 1  }
             
proofstep :: String -> Solver ZProofStep
proofstep name = do
  goal <- sGoal <$> ask
  return (PL.ProofStep name goal)
             
